import numpy as np
from maix import camera,display,image,uart,app

class KalmanFilter:
    def __init__(self, dt=1.0):
        """ 初始化卡尔曼滤波器 """
        # 状态维度：8 (x, y, w, h, vx, vy, vw, vh)
        self.n = 8
        # 观测维度：4 (x, y, w, h)
        self.m = 4
        # 跟踪器是否已初始化
        self.initialized = False
        
        # 状态转移矩阵 F (8x8)
        self.F = np.eye(self.n)
        self.F[0, 4] = dt  # x 位置受 vx 影响
        self.F[1, 5] = dt  # y 位置受 vy 影响
        self.F[2, 6] = dt  # w 尺寸受 vw 影响
        self.F[3, 7] = dt  # h 尺寸受 vh 影响
        
        # 观测矩阵 H (4x8) - 只能观测位置和尺寸
        self.H = np.eye(self.m, self.n)
        
        # 过程噪声协方差 Q - 调整这些值以适应您的场景
        self.Q = np.diag([0.7, 0.7, 0.07, 0.07, 0.8, 0.8, 0.08, 0.08]) ** 2
        
        # 观测噪声协方差 R - 反映检测精度
        self.R = np.diag([1.0, 1.0, 0.5, 0.5]) ** 2
        
        # 状态向量和协方差矩阵
        self.x = np.zeros((self.n, 1))  # 状态 [x, y, w, h, vx, vy, vw, vh]
        self.P = np.eye(self.n) * 100  # 初始高不确定性
    
    def init(self, z):
        """ 使用第一个观测值初始化跟踪器 """
        z = np.array(z).reshape((self.m, 1))
        self.x[:4] = z  # 初始化位置和尺寸
        self.initialized = True 
    
    def predict(self):
        """ 预测步骤：计算先验估计 """
        # 预测状态：x_prior = F * x
        self.x = self.F @ self.x
        
        # 预测协方差：P_prior = F * P * F^T + Q
        self.P = self.F @ self.P @ self.F.T + self.Q
        
        # 返回先验估计（预测值）
        return self.x.copy(), self.P.copy()
    
    def update(self, z):
        """
        更新步骤：结合观测值计算后验估计
        :param z: 观测向量 [x, y, w, h] (4x1)
        """
        # 如果未初始化，先初始化
        if not self.initialized:
            self.init(z)
            return 

        z = np.array(z).reshape((self.m, 1))
        
        # 计算残差：y = z - H * x_prior
        y = z - self.H @ self.x
        
        # 残差协方差：S = H * P_prior * H^T + R
        S = self.H @ self.P @ self.H.T + self.R
        
        # 卡尔曼增益：K = P_prior * H^T * S^{-1}
        K = self.P @ self.H.T @ np.linalg.inv(S)
        
        # 更新状态：x_posterior = x_prior + K * y
        self.x = self.x + K @ y
        
        # 更新协方差：P_posterior = (I - K * H) * P_prior
        I = np.eye(self.n)
        self.P = (I - K @ self.H) @ self.P

        return 
    
    def get_state(self):
        """ 返回当前状态向量 """
        return self.x.flatten()
    
    def get_bbox(self):
        if not self.initialized:
            return [0, 0, 0, 0]  # 防止未初始化时出错
        return self.x[:4].flatten().astype(int)





if __name__ == "__main__":
    cam = camera.Camera()
    dis = display.Display()
    kf = KalmanFilter(dt=0.033)  
    threshold = [[50,70,40,60,-10,10]]
    while not app.need_exit(): 
        img = cam.read()

        # 先进行预测（如果已初始化）
        if kf.initialized:
            kf.predict()   
        
        #获得卡尔曼的观测值
        blobs = img.find_blobs(threshold)
        if blobs:
            blob = max(blobs , key = lambda blob:blob.area())
            x = blob.x()
            y = blob.y()
            w = blob.w()
            h = blob.h()
            img.draw_rect(x,y,w,h,image.COLOR_RED)
            # 使用观测值更新卡尔曼滤波器
            kf.update([x, y, w, h])
            # 提取平滑后的边界框
            e_x, e_y, e_w, e_h = kf.get_bbox()
            img.draw_rect(e_x,e_y,e_w,e_h,image.COLOR_GREEN)
        else:
            if kf.initialized:
                e_x,e_y,e_w,e_h = kf.get_bbox()
                img.draw_rect(e_x,e_y,e_w,e_h,image.COLOR_BLUE)

        dis.show(img)


